#
# example022.py
#
# The wavelet transform
#
# Copyright (C) 2012 Robert Buj Gelonch
# Copyright (C) 2012 David Megias Jimenez
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
__author__ = "Robert Buj Gelonch, and David Megias Jimenez"
__copyright__ = "Copyright 2012, Robert Buj Gelonch and David Megias Jimenez"
__credits__ = ["Robert Buj Gelonch", "David Megias Jimenez"]
__license__ = "GPL"
__version__ = "3"
__maintainer__ = "Robert Buj"
__email__ = "rbuj@uoc.edu"
__status__ = "Development"
__docformat__ = 'plaintext'



from numpy import abs
from numpy import arange
from numpy import array
from numpy import sinc
from numpy import zeros
from pylab import figure
from pylab import plot
from pylab import show
from pylab import stem
from pylab import subplot
from pylab import title
from pywt import Wavelet
from pywt import dwt
from pywt import idwt
from scipy.signal import convolve
from scipy.signal import daub
from scipy.signal import fftconvolve
from scipy.signal import freqz
from scipy.signal import qmf

print "example022.py"
print
print "The discrete wavelet transform"
print
#--------------------------------------------------------------
# Signsl
#--------------------------------------------------------------
x = sinc(arange(-20., 21.) / 5)
#--------------------------------------------------------------
# Wavelet Transform
#--------------------------------------------------------------
# zpd (zero-padding): signal is extended by adding zero samples
# cpd (constant-padding): border values are replicated
# sym (symmetric-padding): signal is extended by mirroring samples
# ppd (periodic-padding): signal is treated as a periodic one
# sp1 (smooth-padding): signal is extended according to the first
#      derivatives calculated on the edges (straight line)
# per (periodization): is like periodic-padding but gives the
#     smallest possible number of decomposition coefficients.
#     IDWT must be performed with the same mode
# http://www.pybytes.com/pywavelets/ref/signal-extension-modes.html#modes
# http://wavelets.pybytes.com/
#--------------------------------------------------------------
mode = 'zpd'
wavelet1 = Wavelet('db10')
cA, cD = dwt(x, wavelet1, mode)

fig = figure(num=0, figsize=(14, 9), dpi=80, facecolor='w', edgecolor='k')
fig.suptitle("DWT & IDWT")
ax2 = subplot(111)
ax2.set_xlim([0, len(x)])
plot(idwt(cA, None, wavelet1, mode) + idwt(None, cD, wavelet1, mode='sym'), \
     'bo-', \
     label="IDWT(cA,None, db2)+IDWT(None, cD, db2)")
plot(idwt(cA, None, wavelet1, mode), \
     'ro--', \
     label="IDWT(cA,None, db2)")
plot(idwt(None, cD, wavelet1, mode), \
     'go--', \
     label="IDWT(None, cD, db2)")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)
fig.savefig("../plots/example022-py-1.eps", format="eps")
show()

################################################
# Filter coefficients / impulse response
################################################
# http://www.pybytes.com/pywavelets/ref/wavelets.html
fig = figure(num=None, \
             figsize=(14, 9), \
             dpi=80, \
             facecolor='w', \
             edgecolor='k')
fig.suptitle("Filter coefficients")
ax2 = subplot(221)
title(r'Decomposition LP filter')
stem(range(len(wavelet1.dec_lo)), wavelet1.dec_lo)
ax2 = subplot(222)
title(r'Decomposition HP filter')
stem(range(len(wavelet1.dec_hi)), wavelet1.dec_hi)
ax2 = subplot(223)
title(r'Reconstruction LP filter')
stem(range(len(wavelet1.rec_lo)), wavelet1.rec_lo)
ax2 = subplot(224)
title(r'Reconstruction HP filter')
stem(range(len(wavelet1.rec_hi)), wavelet1.rec_hi)
fig.savefig("../plots/example022-py-2.eps", format="eps")

################################################
# Filter Freq response
################################################
# http://www.pybytes.com/pywavelets/ref/wavelets.html
fig = figure(num=None, \
             figsize=(9, 14), \
             dpi=80, \
             facecolor='w', \
             edgecolor='k')
fig.suptitle("frequency response of the filters")

w, h = freqz(wavelet1.dec_lo, worN=10 ** 3, whole=False)
subplot(421)
title(r'Decomposition LP filter: H')
plot(abs(h))
subplot(423)
title(r'Decomposition LP filter')
plot(w, h)

w, h = freqz(wavelet1.dec_hi, worN=10 ** 3, whole=False)
subplot(422)
title(r'Decomposition HP filter: H')
plot(abs(h)) # plotting the signal
subplot(424)
title(r'Decomposition HP filter')
plot(w, h)

subplot(425)
title(r'Reconstruction LP filter: H')
w, h = freqz(wavelet1.rec_lo, worN=10 ** 3, whole=False)
plot(abs(h)) # plotting the signal
subplot(427)
title(r'Reconstruction LP filter')
plot(w, h)

w, h = freqz(wavelet1.rec_hi, worN=10 ** 3, whole=False)
subplot(426)
title(r'Reconstruction HP filter: H')
plot(abs(h)) # plotting the signal
subplot(428)
title(r'Reconstruction HP filter')
plot(w, h)
fig.savefig("../plots/example022-py-3.eps", format="eps")

################################################
# wavelet and scaling functions
################################################
phi_1, psi_1, x_1 = wavelet1.wavefun(level=1)
phi_2, psi_2, x_2 = wavelet1.wavefun(level=2)
phi_3, psi_3, x_3 = wavelet1.wavefun(level=3)
phi_4, psi_4, x_4 = wavelet1.wavefun(level=4)

fig = figure(num=None, \
             figsize=(14, 9), \
             dpi=80, \
             facecolor='w', \
             edgecolor='k')
fig.suptitle("wavelet and scaling functions")

ax2 = subplot(211)
title("psi")
plot(x_1, psi_1, 'r--', label="level 1")
plot(x_2, psi_2, 'g--', label="level 2")
plot(x_3, psi_3, 'b-', label="level 3")
plot(x_4, psi_4, 'y.-', label="level 4")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

ax2 = subplot(212)
title("phi")
plot(x_1, phi_1, 'r--', label="level 1")
plot(x_2, phi_2, 'g--', label="level 2")
plot(x_3, phi_3, 'b-', label="level 3")
plot(x_4, phi_4, 'y.-', label="level 4")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)
fig.savefig("../plots/example022-py-4.eps", format="eps")

################################################
# time-domain filtering
################################################
convolve_dec_low = convolve(x, wavelet1.dec_lo, mode='same')
convolve_dec_hight = convolve(x, wavelet1.dec_hi, mode='same')
cA = [convolve_dec_low[2 * i] for i in range(len(convolve_dec_low) / 2)]
cD = [convolve_dec_hight[2 * i] for i in range(len(convolve_dec_hight) / 2)]
A = zeros(len(cA) * 2)
D = zeros(len(cD) * 2)
for i in range(len(cA)):
    A[2 * i] = cA[i]
for i in range(len(cD)):
    D[2 * i] = cD[i]
s = convolve(A, wavelet1.rec_lo, mode='same') + convolve(D, wavelet1.rec_hi, mode='same')

fig = figure(num=4, figsize=(14, 9), dpi=80, facecolor='w', edgecolor='k')
fig.suptitle("DWT Time Domain")

ax2 = subplot(321)
title("Original signal")
plot(x, 'go--', label="x")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

ax2 = subplot(323)
title("Descomposition: convolution")
plot(convolve_dec_low, 'ro--', label="(x*g)")
plot(convolve_dec_hight, 'bo--', label="(x*h)")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

ax2 = subplot(325)
title("Descomposition: Downsample")
plot(cA, 'ro--', label="(x*g)" + unichr(8595) + "2")
plot(cD, 'bo--', label="(x*h)" + unichr(8595) + "2")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

ax2 = subplot(322)
title("Recomposition: upsample & convolution")
plot(s, 'ro--', label="(x*g)" + unichr(8593) + "2")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

fig.savefig("../plots/example022-py-5.eps", format="eps")

################################################
# freq-domain filtering
################################################
convolve_dec_low = fftconvolve(array(x), array(wavelet1.dec_lo), mode='same')
convolve_dec_hight = fftconvolve(array(x), array(wavelet1.dec_hi), mode='same')
cA = [convolve_dec_low[2 * i] for i in range(len(convolve_dec_low) / 2)]
cD = [convolve_dec_hight[2 * i] for i in range(len(convolve_dec_hight) / 2)]
A = zeros(len(cA) * 2)
D = zeros(len(cD) * 2)
for i in range(len(cA)):
    A[2 * i] = cA[i]
for i in range(len(cD)):
    D[2 * i] = cD[i]
s = fftconvolve(A, array(wavelet1.rec_lo), mode='same') + fftconvolve(D, array(wavelet1.rec_hi), mode='same')

fig = figure(num=5, figsize=(14, 9), dpi=80, facecolor='w', edgecolor='k')
fig.suptitle("DWT Time Domain using FFT")

ax2 = subplot(321)
title("Original signal")
plot(x, 'go--', label="x")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

ax2 = subplot(323)
title("Descomposition: convolution")
plot(convolve_dec_low, 'ro--', label="(x*g)")
plot(convolve_dec_hight, 'bo--', label="(x*h)")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

ax2 = subplot(325)
title("Descomposition: Downsample")
plot(cA, 'ro--', label="(x*g)" + unichr(8595) + "2")
plot(cD, 'bo--', label="(x*h)" + unichr(8595) + "2")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

ax2 = subplot(322)
title("Recomposition: upsample & convolution")
plot(s, 'ro--', label="(x*g)" + unichr(8593) + "2")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

fig.savefig("../plots/example022-py-6.eps", format="eps")
show()

lpf_coeff = daub(10)
hpf_coeff = qmf(lpf_coeff)

fig = figure(num=7, figsize=(14, 9), dpi=80, facecolor='w', edgecolor='k')
fig.suptitle("Filter DWT")

subplot(221)
title("LP coeffs")
plot(lpf_coeff, 'bo--', label="x")

subplot(222)
title("HP coeffs")
plot(hpf_coeff, 'bo--', label="x")

fig.savefig("../plots/example022-py-7.eps", format="eps")
show()

print "Done"